"""
数据增强模块
定义各种数据增强方法
"""
import numpy as np
import random
from typing import Callable, List, Optional


class Compose:
    """组合多个变换"""
    def __init__(self, transforms: List[Callable]):
        self.transforms = transforms
    
    def __call__(self, img):
        for t in self.transforms:
            img = t(img)
        return img


class RandomHorizontalFlip:
    """随机水平翻转"""
    def __init__(self, p=0.5):
        self.p = p
    
    def __call__(self, img):
        if random.random() < self.p:
            return np.fliplr(img).copy()
        return img


class RandomVerticalFlip:
    """随机垂直翻转"""
    def __init__(self, p=0.5):
        self.p = p
    
    def __call__(self, img):
        if random.random() < self.p:
            return np.flipud(img).copy()
        return img


class RandomRotation:
    """随机旋转（小角度）"""
    def __init__(self, degrees=5):
        self.degrees = degrees
    
    def __call__(self, img):
        if random.random() < 0.5:
            angle = random.uniform(-self.degrees, self.degrees)
            # 简单实现：只支持90度的倍数
            if abs(angle) > 45:
                k = 1 if angle > 0 else -1
                return np.rot90(img, k).copy()
        return img


class RandomBrightness:
    """随机亮度调整"""
    def __init__(self, brightness_range=(0.8, 1.2)):
        self.min_brightness = brightness_range[0]
        self.max_brightness = brightness_range[1]
    
    def __call__(self, img):
        factor = random.uniform(self.min_brightness, self.max_brightness)
        img = img * factor
        return np.clip(img, 0, 1).astype(np.float32)


class RandomContrast:
    """随机对比度调整"""
    def __init__(self, contrast_range=(0.8, 1.2)):
        self.min_contrast = contrast_range[0]
        self.max_contrast = contrast_range[1]
    
    def __call__(self, img):
        factor = random.uniform(self.min_contrast, self.max_contrast)
        mean = img.mean()
        img = (img - mean) * factor + mean
        return np.clip(img, 0, 1).astype(np.float32)


class RandomNoise:
    """添加随机噪声"""
    def __init__(self, noise_level=0.02):
        self.noise_level = noise_level
    
    def __call__(self, img):
        if random.random() < 0.5:
            noise = np.random.randn(*img.shape) * self.noise_level
            img = img + noise
            return np.clip(img, 0, 1).astype(np.float32)
        return img


class RandomGaussianBlur:
    """随机高斯模糊"""
    def __init__(self, kernel_size=3, sigma_range=(0.1, 2.0)):
        self.kernel_size = kernel_size
        self.sigma_range = sigma_range
    
    def __call__(self, img):
        if random.random() < 0.3:
            from scipy.ndimage import gaussian_filter
            sigma = random.uniform(*self.sigma_range)
            return gaussian_filter(img, sigma=sigma).astype(np.float32)
        return img


class RandomCrop:
    """随机裁剪"""
    def __init__(self, crop_ratio=0.9):
        self.crop_ratio = crop_ratio
    
    def __call__(self, img):
        if random.random() < 0.3:
            h, w = img.shape
            new_h = int(h * self.crop_ratio)
            new_w = int(w * self.crop_ratio)
            
            top = random.randint(0, h - new_h)
            left = random.randint(0, w - new_w)
            
            cropped = img[top:top+new_h, left:left+new_w]
            # 调整回原始大小
            from PIL import Image
            cropped_pil = Image.fromarray((cropped * 255).astype(np.uint8))
            resized = cropped_pil.resize((w, h), Image.BILINEAR)
            return np.array(resized).astype(np.float32) / 255.0
        return img


class Normalize:
    """归一化"""
    def __init__(self, mean=0.5, std=0.5):
        self.mean = mean
        self.std = std
    
    def __call__(self, img):
        return (img - self.mean) / self.std


def get_train_transforms():
    """
    获取训练集数据增强变换
    
    返回:
        transform: 变换组合
    """
    return Compose([
        RandomHorizontalFlip(p=0.5),
        RandomBrightness(brightness_range=(0.85, 1.15)),
        RandomContrast(contrast_range=(0.85, 1.15)),
        RandomNoise(noise_level=0.01),
    ])


def get_val_transforms():
    """
    获取验证集变换（不进行数据增强）
    
    返回:
        transform: 变换组合（空）
    """
    return Compose([])


def get_strong_transforms():
    """
    获取强数据增强变换（用于更大的模型）
    
    返回:
        transform: 变换组合
    """
    return Compose([
        RandomHorizontalFlip(p=0.5),
        RandomVerticalFlip(p=0.3),
        RandomBrightness(brightness_range=(0.8, 1.2)),
        RandomContrast(contrast_range=(0.8, 1.2)),
        RandomNoise(noise_level=0.02),
    ])


if __name__ == "__main__":
    # 测试数据增强
    print("Testing data transforms...")
    
    # 创建测试图像
    test_img = np.random.rand(96, 288).astype(np.float32)
    print(f"Original image shape: {test_img.shape}")
    print(f"Original image range: [{test_img.min():.3f}, {test_img.max():.3f}]")
    
    # 测试训练变换
    train_transform = get_train_transforms()
    augmented_img = train_transform(test_img.copy())
    print(f"\nAugmented image shape: {augmented_img.shape}")
    print(f"Augmented image range: [{augmented_img.min():.3f}, {augmented_img.max():.3f}]")
    
    # 测试多次增强
    print("\nTesting multiple augmentations...")
    for i in range(5):
        aug_img = train_transform(test_img.copy())
        print(f"  Aug {i+1}: mean={aug_img.mean():.3f}, std={aug_img.std():.3f}")
    
    print("\n✓ Transform test passed!")